"""
mark_fire_lane.py
在视频帧上点击标记消防通道区域，自动获取坐标
"""
import cv2
import json


class FireLaneMarker:
    def __init__(self, video_path):
        self.video_path = video_path
        self.points = []
        self.img = None
        self.original_img = None
        self.window_name = 'Mark Fire Lane - 点击标记消防通道区域'

    def mouse_callback(self, event, x, y, flags, param):
        """鼠标回调函数"""
        if event == cv2.EVENT_LBUTTONDOWN:
            # 左键点击：添加点
            self.points.append([x, y])
            print(f"✅ 点 {len(self.points)}: ({x}, {y})")
            self.redraw()

        elif event == cv2.EVENT_RBUTTONDOWN:
            # 右键点击：撤销上一个点
            if self.points:
                removed = self.points.pop()
                print(f"❌ 撤销点: ({removed[0]}, {removed[1]})")
                self.redraw()

        elif event == cv2.EVENT_MOUSEMOVE:
            # 鼠标移动：显示坐标
            temp_img = self.img.copy()
            cv2.putText(temp_img, f"({x}, {y})", (x + 10, y - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1)
            cv2.imshow(self.window_name, temp_img)

    def redraw(self):
        """重新绘制"""
        self.img = self.original_img.copy()

        # 绘制所有点
        for i, point in enumerate(self.points):
            # 绘制点
            cv2.circle(self.img, tuple(point), 8, (0, 255, 0), -1)
            # 绘制编号
            cv2.putText(self.img, str(i + 1), (point[0] + 15, point[1] - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)

        # 绘制线段
        if len(self.points) > 1:
            for i in range(len(self.points) - 1):
                cv2.line(self.img, tuple(self.points[i]),
                         tuple(self.points[i + 1]), (0, 255, 0), 3)

        # 如果点数>=3，绘制闭合线段
        if len(self.points) >= 3:
            cv2.line(self.img, tuple(self.points[-1]),
                     tuple(self.points[0]), (0, 255, 255), 3)

            # 填充半透明多边形
            overlay = self.img.copy()
            import numpy as np
            pts = np.array(self.points, np.int32)
            cv2.fillPoly(overlay, [pts], (0, 255, 0))
            cv2.addWeighted(overlay, 0.3, self.img, 0.7, 0, self.img)

        # 显示提示信息
        info_y = 30
        cv2.putText(self.img, f"Points: {len(self.points)}", (10, info_y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        cv2.putText(self.img, "Left Click: Add | Right Click: Undo | Q: Finish | R: Reset",
                    (10, info_y + 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

        cv2.imshow(self.window_name, self.img)

    def mark(self):
        """开始标记"""
        # 读取视频第一帧
        cap = cv2.VideoCapture(self.video_path)
        if not cap.isOpened():
            print(f"❌ 无法打开视频: {self.video_path}")
            return None

        ret, frame = cap.read()

        # 获取视频信息
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = cap.get(cv2.CAP_PROP_FPS)

        cap.release()

        if not ret:
            print("❌ 无法读取视频帧")
            return None

        self.original_img = frame.copy()
        self.img = frame.copy()

        print("\n" + "=" * 70)
        print("🎥 视频信息")
        print("=" * 70)
        print(f"文件: {self.video_path}")
        print(f"分辨率: {width} x {height} 像素")
        print(f"帧率: {fps:.2f} FPS")
        print("=" * 70)
        print("\n📍 操作说明:")
        print("  👆 左键点击: 按顺序标记消防通道的顶点（建议4个点形成区域）")
        print("  👆 右键点击: 撤销上一个点")
        print("  ⌨️  按 'R' 键: 重置所有点，重新开始")
        print("  ⌨️  按 'Q' 键: 完成标记，保存坐标")
        print("=" * 70)
        print("\n💡 提示:")
        print("  - 沿着消防通道边界顺时针或逆时针点击")
        print("  - 至少需要3个点，推荐4个点（梯形或矩形）")
        print("  - 鼠标移动时会显示当前坐标")
        print("=" * 70 + "\n")

        # 创建窗口
        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        cv2.resizeWindow(self.window_name, 1280, 720)  # 调整窗口大小便于操作
        cv2.setMouseCallback(self.window_name, self.mouse_callback)

        self.redraw()

        while True:
            key = cv2.waitKey(1) & 0xFF

            if key == ord('q') or key == ord('Q'):
                # 完成标记
                if len(self.points) < 3:
                    print("\n⚠️  至少需要3个点才能形成区域！")
                    continue
                break
            elif key == ord('r') or key == ord('R'):
                # 重置
                print("\n🔄 重置所有点")
                self.points = []
                self.redraw()

        cv2.destroyAllWindows()

        if len(self.points) < 3:
            print("\n❌ 标记取消")
            return None

        # 保存结果
        result = {
            'video_path': self.video_path,
            'video_resolution': {
                'width': width,
                'height': height
            },
            'pixel_position': self.points,
            'points_count': len(self.points)
        }

        # 保存到JSON文件
        output_file = 'fire_lane_position.json'
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(result, f, indent=2, ensure_ascii=False)

        # 打印结果
        print("\n" + "=" * 70)
        print("✅ 标记完成！")
        print("=" * 70)
        print(f"标记点数: {len(self.points)}")
        print(f"\npixel_position = {self.points}")
        print("=" * 70)
        print(f"\n📄 配置已保存到: {output_file}")
        print("\n📋 复制以下代码到你的测试脚本或API调用中:")
        print("-" * 70)
        print(f"pixel_position = {self.points}")
        print("-" * 70)
        print("\n🔧 或用于 API 调用:")
        print("-" * 70)
        print(f'pixel_position = "{self.points}"')
        print("-" * 70)

        return self.points


def main():
    """主函数"""
    # 你的视频路径
    video_path = r"E:\zhihao\2025文档\25年10月\Captures\消防通道\消防通道_1.mp4"

    marker = FireLaneMarker(video_path)
    pixel_position = marker.mark()

    if pixel_position:
        print("\n✨ 下一步:")
        print(f"  1. 使用这个坐标测试检测效果")
        print(f"  2. pixel_position = {pixel_position}")


if __name__ == "__main__":
    main()